from selenium import webdriver
from selenium.webdriver.firefox.service import Service
from selenium.common.exceptions import  TimeoutException
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as ec
from selenium.webdriver.chrome.service import Service
import os
import time
import random
import ddddocr

#存储最后下载的ID和用户信息
LAST_ID_FILE = "last_id.txt"
USER_INFO_FILE = "cnvd_cookie.txt"

# 下载文件夹
DOWNLOAD_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "CNVD_DB")
os.makedirs(DOWNLOAD_DIR, exist_ok=True)

# 获取当前目录路径
sys_path = os.path.abspath(".")

# ChromeDriver路径
CHROMEDRIVER_PATH = os.path.join(sys_path, "chromedriver", "chromedriver")

# CNVD基础URL模板
#BASE_URL = "https://www.cnvd.org.cn/shareData/list"
BASE_URL = "https://www.cnvd.org.cn/user/login"

# OCR工具
ocr = ddddocr.DdddOcr()

def get_last_id():
    if os.path.exists(LAST_ID_FILE):
        with open(LAST_ID_FILE, 'r', encoding='utf-8') as f:
            try:
                return int(f.read().strip())
            except:
                return 1
    else:
        return 1

def save_last_id(cid):
    with open(LAST_ID_FILE, 'w', encoding='utf-8') as f:
        f.write(str(cid))
def get_user_info():
    if os.path.exists(USER_INFO_FILE):
        with open(USER_INFO_FILE, 'r', encoding='utf-8') as f:
            lines = f.readlines()
            if len(lines) >= 2:
                email = lines[0].strip()
                password = lines[1].strip()
                return email, password
    return None, None

def save_user_info(email, password):
    with open(USER_INFO_FILE, 'w', encoding='utf-8') as f:
        f.write(email + '\n')
        f.write(password + '\n')

# 截图并识别验证码
def recognize_captcha(driver):
    # 定位验证码图片元素
    try:
        captcha_img = WebDriverWait(driver, 10).until(
            ec.presence_of_element_located((By.ID, 'codeSpan'))
        )
        # 直接截图该元素
        captcha_img.screenshot('code_png.png')
    except Exception as e:
        print("验证码截图失败:", e)
        driver.save_screenshot('captcha_error.png')
        raise

    # OCR识别
    my_ocr = ddddocr.DdddOcr()
    with open('code_png.png', 'rb') as f:
        code_img_bytes = f.read()
    code = my_ocr.classification(code_img_bytes)

    # 切换标签页为登录页面
    driver.switch_to.window(driver.window_handles[0])
    print(f"识别的验证码: {code}")
    return code

def login(driver, email, password):
    print("正在尝试登录...")
    try:
        email_input = WebDriverWait(driver, 10).until(ec.presence_of_element_located((By.ID, 'email')))
        password_input = driver.find_element(By.ID, 'password')
        code_input = driver.find_element(By.ID, 'myCode')
        time.sleep(2)
        email_input.clear()
        password_input.clear()
        code_input.clear()

        # 输入用户名、密码
        for attempt in range(3):  # 最多重试3次
            # 重新获取一下元素
            email_input = driver.find_element(By.ID, 'email')
            password_input = driver.find_element(By.ID, 'password')
            code_input = driver.find_element(By.ID, 'myCode')
            time.sleep(2)
            email_input.clear()
            password_input.clear()
            code_input.clear()

            for char in email:
                email_input.send_keys(char)
                time.sleep(random.uniform(0.05, 0.15))
            for char in password:
                password_input.send_keys(char)
                time.sleep(random.uniform(0.05, 0.15))
            time.sleep(2)  # 等待2秒后重试

            # 识别验证码
            code = recognize_captcha(driver)
            for char in code:
                code_input.send_keys(char)
                time.sleep(random.uniform(0.05, 0.15))

            # 提交登录
            login_button = driver.find_element(By.XPATH, '//a[@class="btn b_28_1"]')
            login_button.click()

            time.sleep(3)
            if "login" not in driver.current_url:
                print(" 登录成功！")
                return True

            print(f"登录失败，正在进行第{attempt + 1}次重试...")
            time.sleep(2)

    except Exception as e:
        print("登录过程中出错：", e)
        return False


def download_latest_file(driver, last_id):
    try:
        print("开始查找最新漏洞链接...")
        # 定位第一条记录的链接
        first_link = WebDriverWait(driver, 10).until(
            ec.element_to_be_clickable((By.XPATH, '//table[@class="tlist"]/tbody/tr[1]/td[1]/a'))
        )
        file_id = first_link.get_attribute('href').split("/")[-1]
        print(f"获取到漏洞ID: {file_id}")

        # 判断是否为新文件
        if int(file_id) <= last_id:
            print(f"当前最新漏洞ID {file_id}，小于等于上次已下载ID {last_id}，无需重复下载。")
            return last_id

        # 查找并点击下载链接
        print("等待下载按钮出现...")
        download_link = WebDriverWait(driver, 10).until(
            ec.element_to_be_clickable((By.XPATH, '//a[@title="下载xml"]'))
        )

        # 再次确认下载链接对应的 ID 是否一致
        download_file_id = download_link.get_attribute('href').split("/")[-1]
        if download_file_id == file_id:
            download_link.click()
            print(f"成功触发下载漏洞文件: {file_id}.xml")

            # 更新 last_id
            save_last_id(file_id)
            print(f"已更新 last_id 为: {file_id}")
            return file_id
        else:
            print("下载链接与当前漏洞ID不匹配，跳过下载。")
            return last_id

    except Exception as e:
        print("下载文件时发生错误：", e)
        return last_id

def delete_code_png():
    if os.path.exists('code_png.png'):
        os.remove('code_png.png')
        print("已删除 code_png.png")
    else:
        print("code_png.png 不存在")

def main():
    options = webdriver.FirefoxOptions()
    #options.add_argument('--headless')
    options.add_argument('--no-sandbox')
    options.add_argument('--disable-dev-shm-usage')
    #options.add_argument('--disable-gpu')

    # 可选：设置窗口大小
    options.add_argument("--window-size=1920,1080")

    service = Service(executable_path='/usr/bin/geckodriver')
    driver = webdriver.Firefox(service=service, options=options)

    driver.execute_script("Object.defineProperty(navigator, 'webdriver', {get: () => undefined})")

    # driver.get(BASE_URL)
    # print(driver.title)
    # driver.quit()

    try:
        current_last_id = get_last_id()
        print(f"当前记录的最后下载漏洞ID: {current_last_id}")

        # 获取缓存的账号密码
        email, password = get_user_info()
        if not email or not password:
            email = input("请输入CNVD账号邮箱：").strip()
            password = input("请输入CNVD账号密码：").strip()
            save_user_info(email, password)

        # 进入CNVD列表页
        print(f"正在访问URL: {BASE_URL}")
        driver.get(BASE_URL)

        # 检查是否需要登录
        try:
            WebDriverWait(driver, 5).until(
                ec.url_contains("login")
            )
            print("检测到登录页面，开始登录流程...")
            if not login(driver, email, password):
                raise Exception("登录失败，无法继续操作")
        except TimeoutException:
            print("当前已处于登录状态")

        # 开始下载
        new_last_id = download_latest_file(driver, current_last_id)
        print(f" 当前最新漏洞已下载完成，更新最后ID为: {new_last_id}")

    finally:
        time.sleep(3)  # 添加短暂停顿，确保浏览器已关闭
        delete_code_png()
        driver.quit()

if __name__ == "__main__":
    main()